import os
import random
import numpy as np
import scipy.io
import matplotlib.pyplot as plt
from scipy.signal import medfilt
from scipy.interpolate import interp1d


def delt(y: np.array) -> np.array:
    """
    计算一阶差分
    """
    return np.diff(y, prepend=y[0])

def detect_transition_regions(y: np.array, T=8.0*0.8, kernel_size=51, min_region_size=100):
    """
    识别突变区域
    :param y: 输入信号
    :param T: 梯度变化阈值
    :param kernel_size: 中值滤波窗口大小
    :param min_region_size: 认为是有效突变区域的最小长度
    :return: 突变区域的索引列表 [(start1, end1), (start2, end2), ...]
    """
    kernel_size = kernel_size if kernel_size % 2 == 1 else kernel_size + 1

    y_delt = delt(y)
    y_delt_filtered = medfilt(np.abs(y_delt), kernel_size=kernel_size)  # 过滤掉孤立噪声点
    jumps = np.where(y_delt_filtered > T)[0]  # 突变点索引

    transition_regions = []
    start = None
    for i in range(len(jumps) - 1):
        if start is None:
            start = jumps[i]
        if jumps[i+1] - jumps[i] > 1:  # 断开，则标记为一个区域
            if (jumps[i] - start) > min_region_size:
                transition_regions.append((start, jumps[i]))
            start = None
    
    if start is not None and (jumps[-1] - start) > min_region_size:
        transition_regions.append((start, jumps[-1]))

    return transition_regions

# def repair_signal(y: np.array, transition_regions, method='linear'):
#     """
#     修复突变区域
#     :param y: 原始信号
#     :param transition_regions: 突变区域索引列表
#     :param method: 插值方法 ('linear', 'cubic', 'spline')
#     :return: 修复后的信号
#     """
#     y_fixed = y.copy()
    
#     for start, end in transition_regions:
#         # 取突变区域前后的数据点
#         x = [start - 1, end + 1]
#         y_vals = [y_fixed[start - 1], y_fixed[end + 1]]
        
#         # 进行插值填充
#         f = interp1d(x, y_vals, kind=method, fill_value="extrapolate")
#         y_fixed[start:end+1] = f(np.arange(start, end+1))
    
#     return y_fixed


def repair_signal(y: np.array, transition_regions):
    """
    修复突变区域，并拼接上下限
    :param y: 原始信号
    :param transition_regions: 突变区域索引列表 [(start, end), ...]
    :return: 修复后的信号
    """
    y_fixed = np.copy(y)

    for start, end in transition_regions:
        if start == 0 or end == len(y) - 1:
            continue  # 避免越界
        
        # 取突变区域前后的数据点
        y_before = y_fixed[start - 1]
        y_after = y_fixed[end + 1]

        # 上限 → 突变区 → 下限
        if y_before > y_after:
            correction = y_before - y_after  # 计算突变区域的变化量
            y_fixed[start:end+1] -= correction  # 直接补偿

        # 下限 → 突变区 → 上限
        elif y_before < y_after:
            correction = y_after - y_before
            y_fixed[start:end+1] += correction  # 直接补偿

    return y_fixed


def deal_hcn_data(file_name: str):
    data = scipy.io.loadmat(file_name)
    period = 8  # 设定周期
    indexs = random.sample(range(1, 51), 4)  # 随机抽取4组数据

    for ii in indexs:
        y_original = data[f'fai{ii}'].flatten()
        y_true = data[f'ne{ii}'].flatten()

        # 计算突变区域
        transition_regions = detect_transition_regions(y_original)
        # y_repaired = repair_signal(y_original, transition_regions, method='linear')
        y_repaired = repair_signal(y_original, transition_regions)

        plt.figure(figsize=(12, 9))

        # 原始信号
        plt.subplot(3, 1, 1)
        plt.scatter(range(len(y_original)), y_original, label=f'fai{ii} (original)', s=1)
        plt.title('Original Signal')
        plt.legend()

        # 差分信号
        plt.subplot(3, 1, 2)
        y_delt = delt(y_original)
        plt.plot(y_delt, label=f'y_delt {ii}')
        plt.title('Difference Signal')
        plt.legend()

        # 修复后信号
        plt.subplot(3, 1, 3)
        plt.plot(y_true, label=f'ne{ii} (true)')
        plt.plot(y_repaired, label=f'ne{ii} (repaired)', linestyle="--")
        
        # 标注突变区域
        for start, end in transition_regions:
            plt.axvspan(start, end, color='red', alpha=0.3, label="Jump Region" if start == transition_regions[0][0] else "")
        
        plt.title('Repaired Signal')
        plt.legend()
        
        plt.tight_layout()
        plt.show()

if __name__ == '__main__':
    file_name = 'sjj1.mat'
    deal_hcn_data(file_name)
    print('Done!')
